import torch
from torch import nn
from torch.nn import functional as F


IMG_SIZE = 224
LR = 0.001


class InceptionBlock(nn.Module):
    def __init__(self, in_channels, c1, c2, c3, c4):
        super().__init__()
        self.s1 = nn.Sequential(
            nn.Conv2d(in_channels, c1[0], kernel_size=1), nn.ReLU()
        )
        self.s2 = nn.Sequential (
            nn.Conv2d(in_channels, c2[0], kernel_size=1), nn.ReLU(),
            nn.Conv2d(c2[0], c2[1], kernel_size=3, padding=1), nn.ReLU()
        )
        self.s3 = nn.Sequential (
            nn.Conv2d(in_channels, c3[0], kernel_size=1), nn.ReLU(),
            nn.Conv2d(c3[0], c3[1], kernel_size=5, padding=2), nn.ReLU(),
        )
        self.s4 = nn.Sequential (
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_channels, c4[0], kernel_size=1), nn.ReLU()
        )

    
    def forward(self, input):
        p1 = self.s1(input)
        p2 = self.s2(input) 
        p3 = self.s3(input) 
        p4 = self.s4(input) 
        return torch.cat((p1, p2, p3, p4), dim=1)
    

class AuxiliaryClassifier(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.AvgPool2d(kernel_size=5, stride=3),
            nn.Conv2d(in_channels, 128, kernel_size=1, stride=1), nn.ReLU(),
            nn.Flatten(),
            nn.Linear(128 * 4 * 4, 1024),
            nn.Dropout(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        return self.net(input)


class InceptionV1(nn.Module):
    def __init__(self, num_classes, aux_logits=True, in_channels=3):
        super(InceptionV1, self).__init__()
        self.aux_logits = aux_logits
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 192, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(192), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self.inc3 = nn.Sequential(
            InceptionBlock(192, (64,), (96, 128), (16, 32), (32,)),
            InceptionBlock(256, (128,), (128, 192), (32, 96), (64,)),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )
        self.a4 = InceptionBlock(480, (192,), (96, 208), (16, 48), (64,))
        self.b4 = InceptionBlock(512, (160,), (112, 224), (24, 64), (64,))
        self.b4_aux = AuxiliaryClassifier(512, num_classes)
        self.c4 = InceptionBlock(512, (128,), (128, 256), (24, 64), (64,))
        self.d4 = InceptionBlock(512, (112,), (144, 288), (32, 64), (64,))
        self.e4 = InceptionBlock(528, (256,), (160, 320), (32, 128), (128,))
        self.e4_aux = AuxiliaryClassifier(528, num_classes)
        self.inc5 = nn.Sequential(
            InceptionBlock(832, (256,), (160, 320), (32, 128), (128,)),
            InceptionBlock(832, (384,), (192, 384), (48, 128), (128,)),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.4),
            nn.Flatten(),
            nn.Linear(1024, num_classes)
        )


    def forward(self, input):
        y = self.conv1(input)
        y = self.conv2(y)
        y = self.inc3(y)
        y_a4 = self.a4(y)
        y = self.b4(y_a4)
        y = self.c4(y)
        y_d4 = self.d4(y)
        y = self.e4(y_d4)
        y = F.max_pool2d(y, kernel_size=3, stride=2, padding=1)
        y = self.inc5(y)
        y = self.classifier(y)

        if (self.aux_logits and self.training):
            y_b4_aux = self.b4_aux(y_a4)
            y_e4_aux = self.e4_aux(y_d4)
            y = y + y_b4_aux * 0.3 + y_e4_aux * 0.3

        return y
        
    

    @property
    def image_size(self):
        return IMG_SIZE
    

    @property
    def lr(self):
        return LR
    